MODULE sefmm_V_driver

   USE fmm_global_paras
   IMPLICIT NONE
   PRIVATE
   ! Public procedures
   PUBLIC :: sefmm_init_potential,        &
             sefmm_free_potential,        &
             sefmm_build_NF_potential,    &
             sefmm_get_potential

   TYPE(sefmm_box_potential), ALLOCATABLE, TARGET, SAVE :: box_pot(:)
   REAL(REALK), SAVE :: used_memory  ! MB

CONTAINS

!-------------------------------------------------------------------------------

   SUBROUTINE sefmm_init_potential(nbox,scheme)

      USE sefmm_cheb_Fij,   ONLY: cheb_init_Fij
      USE sefmm_quadrature, ONLY: sefmm_init_quadrature

      IMPLICIT NONE
      TYPE(scheme_paras), INTENT(IN) :: scheme
      INTEGER(INTK),      INTENT(IN) :: nbox

      INTEGER(INTK) :: i

      CALL cheb_init_Fij(scheme)
      CALL sefmm_init_quadrature(scheme)

      ALLOCATE(box_pot(nbox))
      DO i = 1, nbox 
         box_pot(i)%order = 0
         NULLIFY(box_pot(i)%V)
      END DO
      used_memory = zero

   END SUBROUTINE sefmm_init_potential

!-------------------------------------------------------------------------------

   SUBROUTINE sefmm_free_potential

      USE sefmm_cheb_Fij, ONLY: cheb_free_Fij
      IMPLICIT NONE
      INTEGER(INTK) :: i

      CALL cheb_free_Fij

      DO i = 1, SIZE(box_pot)
         box_pot(i)%order = 0
         IF (ASSOCIATED(box_pot(i)%V)) DEALLOCATE(box_pot(i)%V)
         NULLIFY(box_pot(i)%V)
      END DO
      DEALLOCATE(box_pot)
      used_memory = zero

   END SUBROUTINE sefmm_free_potential

!-------------------------------------------------------------------------------

   SUBROUTINE sefmm_get_potential(box,Vpot,order)

      IMPLICIT NONE
      INTEGER(INTK),        INTENT(IN)  :: box
      REAL(REALK), POINTER, INTENT(OUT) :: Vpot(:,:,:)
      INTEGER(INTK),        INTENT(OUT) :: order

      order = box_pot(box)%order
      IF (order /= 0) THEN
         Vpot => box_pot(box)%V(:,:,:)
      ELSE
         NULLIFY (Vpot)
      END IF

   END SUBROUTINE sefmm_get_potential

!-------------------------------------------------------------------------------

   SUBROUTINE sefmm_build_NF_potential(nbox,grain,NN_factor)

      USE fmm_utils,     ONLY: fmm_loop_progress
      USE sefmm_density, ONLY: sefmm_get_spectral_density
!      USE sefmm_density, ONLY: sefmm_get_spectral_density,     &
!                               sefmm_get_max_NN_order

      IMPLICIT NONE
      INTEGER(INTK), INTENT(IN) :: nbox
      REAL(REALK),   INTENT(IN) :: grain, NN_factor

      REAL(REALK), POINTER :: rho(:,:,:)
      REAL(REALK)   :: T0, T1
      INTEGER(INTK) :: i, order, NN_order, tmp

      CALL CPU_TIME(T0)

      ! Build self-potential
      DO i = 1, nbox
         CALL fmm_loop_progress(i,nbox)
         CALL sefmm_get_spectral_density(i,rho,order)
         IF (order == 0) CYCLE
         CALL sefmm_do_Coulomb_quad(i,0,order,order,grain,rho)
      END DO

      CALL CPU_TIME(T1)
      CALL TIMTXT('>>> TIME USED get self-potential', T1-T0, LUPRI)

      ! Build nearest-neighbour potentials
      IF (nbox == 1) RETURN
      DO i = 1, nbox
         CALL fmm_loop_progress(i,nbox)
         CALL sefmm_get_spectral_density(i,rho,order)
         IF (order == 0) CYCLE
         !CALL sefmm_get_max_NN_order(i,tmp)
         !NN_order = MIN(order,tmp)
         NN_order = INT(order*NN_factor)
         CALL sefmm_do_Coulomb_quad(i,1,order,NN_order,grain,rho)
      END DO

      CALL CPU_TIME(T0)
      CALL TIMTXT('>>> TIME USED get NN-potential', T0-T1, LUPRI)
      WRITE(LUPRI,'(A,F7.2)') " NF-potential memory (MB) = ", used_memory

   END SUBROUTINE sefmm_build_NF_potential

!-------------------------------------------------------------------------------

   SUBROUTINE sefmm_do_Coulomb_quad(box,nnbox,nrho,npot,grain,rho)

      USE sefmm_quadrature, ONLY: get_Coulomb_quadrature
      USE sefmm_cheb_Fij,   ONLY: cheb_get_Fij

      IMPLICIT NONE
      INTEGER(INTK), INTENT(IN) :: box         ! rho-box ID
      INTEGER(INTK), INTENT(IN) :: nnbox       ! width of NN space for V
      INTEGER(INTK), INTENT(IN) :: nrho, npot  ! Cheb order of rho and V
      REAL(REALK),   INTENT(IN) :: grain
      REAL(REALK),   INTENT(IN) :: rho(nrho,nrho,nrho)

      REAL(REALK), POINTER :: Fijt(:,:,:), Fij(:,:)
      REAL(REALK), POINTER :: rts(:), wts(:)
      REAL(REALK)   :: V((1+2*nnbox)*npot,(1+2*nnbox)*npot,(1+2*nnbox)*npot)
      INTEGER(INTK) :: nt

      NULLIFY (rts, wts, Fijt, Fij)

      CALL get_Coulomb_quadrature(nnbox,rts,wts)

      CALL cheb_get_Fij(nnbox,nrho,(1+2*nnbox)*npot,rts,Fijt)

!V(1,1,1) =zero
!DO nt = 1, SIZE(rts)
!   V(1,1,1) = V(1,1,1) + wts(nt)*EXP( -(rts(nt)*1.0d0)**2)
!endDO
!print *, one/1.0d0, v(1,1,1)
!V(1,1,1) =zero
!DO nt = 1, SIZE(rts)
!   V(1,1,1) = V(1,1,1) + wts(nt)*EXP( -(rts(nt)*2.0d0)**2)
!endDO
!print *, one/2.0d0, v(1,1,1)
!V(1,1,1) =zero
!DO nt = 1, SIZE(rts)
!   V(1,1,1) = V(1,1,1) + wts(nt)*EXP( -(rts(nt)*3.0d0)**2)
!endDO
!print *, one/3.0d0, v(1,1,1)
!if (nnbox > 0) stop

      V(:,:,:) = zero
      DO nt = 1, SIZE(rts)
         Fij => Fijt(:,:,nt)
         IF (nnbox == 0) THEN
            CALL sefmm_contract_Fij(nrho,npot,Fij,rho,wts(nt),V)
         ELSE IF (nnbox == 1) THEN
            CALL sefmm_contract_NN_Fij(box,nrho,3*npot,Fij,rho,wts(nt),V)
         END IF
      END DO

      V = V*((half*grain)**2)

      CALL sefmm_save_potential(nnbox,box,(1+2*nnbox)*npot,V)
      NULLIFY (rts, wts, Fijt, Fij)

   END SUBROUTINE sefmm_do_Coulomb_quad

!-------------------------------------------------------------------------------

   SUBROUTINE sefmm_contract_Fij(nrho,npot,Fij,coeffs,weight,Vpot)

      IMPLICIT NONE
      INTEGER(INTK), INTENT(IN)    :: nrho, npot
      REAL(REALK),   INTENT(IN)    :: Fij(nrho,npot)
      REAL(REALK),   INTENT(IN)    :: coeffs(nrho,nrho,nrho)
      REAL(REALK),   INTENT(IN)    :: weight
      REAL(REALK),   INTENT(INOUT) :: Vpot(npot,npot,npot)

      REAL(REALK), ALLOCATABLE :: V(:,:,:)
      REAL(REALK), ALLOCATABLE :: W(:,:,:)
      INTEGER(INTK) :: ix,iy,iz, n

      n = MAX(nrho,npot)
      ALLOCATE (V(n,n,n))
      ALLOCATE (W(n,n,n))

      DO iz = 1, nrho
         DO iy = 1, nrho
            DO ix = 1, npot
               V(ix,iy,iz) = DOT_PRODUCT(coeffs(1:nrho,iy,iz),Fij(1:nrho,ix))
            END DO
         END DO
      END DO

      ! Reorder indices
      DO iz = 1, nrho
         DO ix = 1, npot
            W(1:nrho,ix,iz) = V(ix,1:nrho,iz)
         END DO
      END DO

      DO iz = 1, nrho
         DO ix = 1, npot
            DO iy = 1, npot
               V(iy,ix,iz) = DOT_PRODUCT(W(1:nrho,ix,iz),Fij(1:nrho,iy))
            END DO
         END DO
      END DO

      ! Reorder indices
      DO ix = 1, npot
         DO iy = 1, npot
            W(1:nrho,iy,ix) = V(iy,ix,1:nrho)
         END DO
      END DO

      DO ix = 1, npot
         DO iy = 1, npot
            DO iz = 1, npot
               V(iz,iy,ix) = DOT_PRODUCT(W(1:nrho,iy,ix),Fij(1:nrho,iz))
               Vpot(ix,iy,iz) = Vpot(ix,iy,iz) + weight*V(iz,iy,ix)
            END DO
         END DO
      END DO

      DEALLOCATE (V,W)

   END SUBROUTINE sefmm_contract_Fij

!-------------------------------------------------------------------------------

   SUBROUTINE sefmm_contract_NN_Fij(boxID,nrho,npot,Fij,coeffs,weight,Vpot)

      USE fefmm_box_utils, ONLY: fefmm_get_neighbours

      IMPLICIT NONE
      INTEGER(INTK), INTENT(IN)    :: nrho, npot, boxID
      REAL(REALK),   INTENT(IN)    :: Fij(nrho,npot)
      REAL(REALK),   INTENT(IN)    :: coeffs(nrho,nrho,nrho)
      REAL(REALK),   INTENT(IN)    :: weight
      REAL(REALK),   INTENT(INOUT) :: Vpot(npot,npot,npot)

      REAL(REALK), PARAMETER :: thresh = 1.0d-6
      REAL(REALK), ALLOCATABLE :: V(:,:,:)
      REAL(REALK), ALLOCATABLE :: W(:,:,:)
      INTEGER(INTK) :: hi(3), lo(3)
      INTEGER(INTK) :: ix,iy,iz, n

      n = MAX(nrho,npot)
      ALLOCATE (V(n,n,n))
      ALLOCATE (W(n,n,n))
      V = zero
      W = zero

      CALL fefmm_get_neighbours(boxID,npot/3,hi,lo)

      DO ix = lo(1), hi(1)
         IF ( ABS(Fij(1,ix)) < thresh ) CYCLE
         DO iz = 1, nrho
            DO iy = 1, nrho
               V(iy,iz,ix) = DOT_PRODUCT(coeffs(1:nrho,iy,iz),Fij(1:nrho,ix))
            END DO
         END DO
      END DO

      DO iy = lo(2), hi(2)
         IF ( ABS(Fij(1,iy)) < thresh ) CYCLE
         DO ix = lo(1), hi(1)
            IF ( ABS(Fij(1,ix)) < thresh ) CYCLE
            DO iz = 1, nrho
               W(iz,ix,iy) = DOT_PRODUCT(V(1:nrho,iz,ix),Fij(1:nrho,iy))
            END DO
         END DO
      END DO

      DO iy = lo(1), hi(1)
         IF ( ABS(Fij(1,iy)) < thresh ) CYCLE
         DO ix = lo(2), hi(2)
            IF ( ABS(Fij(1,ix)) < thresh ) CYCLE
            DO iz = lo(3), hi(3)
               IF ( ABS(Fij(1,iz)) < thresh ) CYCLE
               V(iz,ix,iy) = DOT_PRODUCT(W(1:nrho,ix,iy),Fij(1:nrho,iz))
               Vpot(ix,iy,iz) = Vpot(ix,iy,iz) + weight*V(iz,ix,iy)
            END DO
         END DO
      END DO

      DEALLOCATE (V,W)

   END SUBROUTINE sefmm_contract_NN_Fij

!-------------------------------------------------------------------------------

   SUBROUTINE sefmm_save_potential(nnbox,box,npot,Vpot)

      IMPLICIT NONE
      INTEGER(INTK),       INTENT(IN) :: box, nnbox, npot
      REAL(REALK), TARGET, INTENT(IN) :: Vpot(npot,npot,npot)

      SELECT CASE (nnbox)
         CASE (0)
            CALL sefmm_save_self_potential(box,npot,Vpot)
         CASE (1)
            CALL sefmm_save_NN_potential(box,npot,Vpot)
         CASE DEFAULT
            CALL fmm_quit('sefmm_save_potential')
      END SELECT

   END SUBROUTINE sefmm_save_potential

!-------------------------------------------------------------------------------

   SUBROUTINE sefmm_save_self_potential(boxID,npot,Vpot)

      IMPLICIT NONE
      INTEGER(INTK), INTENT(IN) :: boxID, npot
      REAL(REALK),   INTENT(IN) :: Vpot(npot,npot,npot)

      LOGICAL, SAVE :: flag = .TRUE.
      box_pot(boxID)%order = npot

      IF (.NOT. ASSOCIATED (box_pot(boxID)%V)) THEN
         ALLOCATE (box_pot(boxID)%V(npot,npot,npot))
         used_memory = used_memory + (npot**3)*8 /1d6
         IF ((used_memory > 1000) .AND. flag) THEN
            flag = .FALSE.
            WRITE(LUPRI,'(A,F7.2)') " NF-potential memory (MB) = ", used_memory
         END IF
         box_pot(boxID)%V(:,:,:) = zero
      END IF

      box_pot(boxID)%V = box_pot(boxID)%V + Vpot

   END SUBROUTINE sefmm_save_self_potential

!-------------------------------------------------------------------------------

   SUBROUTINE sefmm_translate_and_store(ibox,nrho,Vin,npot)

      USE sefmm_cheb_utils, ONLY: cheb_translate_coeffs

      IMPLICIT NONE
      INTEGER(INTK), INTENT(IN) :: ibox, nrho, npot
      REAL(REALK),   INTENT(IN) :: Vin(npot,npot,npot)

      REAL(REALK), ALLOCATABLE :: Vout(:,:,:)
      LOGICAL, SAVE :: flag = .TRUE.

      IF (.NOT. ASSOCIATED (box_pot(ibox)%V)) THEN
         ALLOCATE (box_pot(ibox)%V(nrho,nrho,nrho))
         used_memory = used_memory + (nrho**3 *8 /1d6)
         IF ((used_memory > 1000) .AND. flag) THEN
            flag = .FALSE.
            WRITE(LUPRI,'(A,F7.2)') " NF-potential memory (MB) = ", used_memory
         END IF
         box_pot(ibox)%V(:,:,:) = zero
      END IF

      ! Translate to different grid quality if needed
      IF (nrho /= npot) THEN
         ALLOCATE (Vout(nrho,nrho,nrho))
         CALL cheb_translate_coeffs(npot,nrho,Vin,Vout)
         box_pot(ibox)%V = box_pot(ibox)%V + Vout
         DEALLOCATE (Vout)
      ELSE
         box_pot(ibox)%V = box_pot(ibox)%V + Vin
      END IF
      box_pot(ibox)%order = nrho

   END SUBROUTINE sefmm_translate_and_store

!-------------------------------------------------------------------------------

   SUBROUTINE sefmm_save_NN_potential(boxID,npot,Vpot)

      USE fefmm_box_utils,  ONLY: fefmm_box_map
      USE sefmm_density,    ONLY: sefmm_get_order

      IMPLICIT NONE
      INTEGER(INTK),       INTENT(IN) :: boxID, npot
      REAL(REALK), TARGET, INTENT(IN) :: Vpot(npot,npot,npot)

      REAL(REALK), POINTER :: V(:,:,:)
      INTEGER(INTK) :: ix,iy,iz, ibox, order, npts
      INTEGER(INTK) :: xlo,xhi, ylo,yhi, zlo,zhi 
      INTEGER(INTK) :: xshift, yshift, zshift

      NULLIFY (V)
      npts = npot/3

      ! Isolate cubes of dimension npts and accumulate to stored V
      DO ix = 1, 3
         xlo = (ix-1)*npts +1
         xhi = xlo + npts -1
         xshift = ix -2
         DO iy = 1, 3
            ylo = (iy-1)*npts +1
            yhi = ylo + npts -1
            yshift = iy -2
            DO iz = 1, 3
               zlo = (iz-1)*npts +1
               zhi = zlo + npts -1
               zshift = iz -2

               ! Skip self-potential
               IF ( (ix==2) .AND.  (iy==2) .AND.  (iz==2) ) CYCLE

               ! Check if neighbour is off the edge of the grid
               ibox = fefmm_box_map(boxID,xshift,yshift,zshift) 
               IF (ibox < 0) CYCLE

               CALL sefmm_get_order(ibox,order)
               box_pot(ibox)%order = order
               ! Don't store potential where there is vanishing density
               IF (order == 0) CYCLE

               V => Vpot(xlo:xhi,ylo:yhi,zlo:zhi)
               CALL sefmm_translate_and_store(ibox,order,V,npts)

            END DO
         END DO
      END DO
      NULLIFY (V)

   END SUBROUTINE sefmm_save_NN_potential

!-------------------------------------------------------------------------------

END MODULE sefmm_V_driver
